
#pragma once
#include <eigen3/Eigen/Core>
#include <pcl/common/transforms.h>
#include <pcl/kdtree/kdtree_flann.h>

#include <g2o/core/g2o_core_api.h>
#include <g2o/core/base_vertex.h>
#include <g2o/core/base_unary_edge.h>
#include <g2o/core/block_solver.h>
#include <g2o/core/optimization_algorithm_levenberg.h>
#include <g2o/core/optimization_algorithm_gauss_newton.h>
#include <g2o/core/optimization_algorithm_dogleg.h>
#include <g2o/solvers/dense/linear_solver_dense.h>

#include "../3rdparty/sophus/se3.hpp"

#include "registration_interface.hpp"

namespace TESTICP
{
    class opt_ICP_G2O : public RegistrationInterface
    {

    public:
        EIGEN_MAKE_ALIGNED_OPERATOR_NEW;
        opt_ICP_G2O(const YAML::Node& node);
        ~opt_ICP_G2O();
        void setInputTarget(const CloudPtr& target) override;
        void setInputSource(const CloudPtr& source) override;
        bool align( CloudType& output_ptr, const Eigen::Matrix4f& predict_pose = Eigen::Matrix4f::Identity()) override;

        float getFitnessScore();
        inline Eigen::Matrix4f getFinalTransformation() { return final_pose; }

    private:
        CloudPtr target_ptr, source_ptr;
        Eigen::Matrix4f final_pose;
        int max_iterations;
        float max_coresspoind_dis;
        float trans_eps;
        float euc_fitness_eps;

        pcl::KdTreeFLANN<PointType>::Ptr kdtree_flann;
    };

    class ICPVertex : public g2o::BaseVertex<6, Sophus::SE3d>
    {
    public:
        EIGEN_MAKE_ALIGNED_OPERATOR_NEW

            virtual void setToOriginImpl() override {
            _estimate = Sophus::SE3d();
        }
        virtual void oplusImpl(const double* update) override {
            Eigen::Matrix<double, 6, 1> delta_r;
            delta_r << update[0], update[1], update[2], update[3], update[4], update[5];
            // _estimate = Sophus::SE3d::exp(delta_r) * _estimate;  //  左乘
            _estimate = _estimate * Sophus::SE3d::exp(delta_r);
        }
        virtual bool read(std::istream& in) { }
        virtual bool write(std::ostream& out) const { }
    };

    class ICPEdge : public g2o::BaseUnaryEdge<3, Eigen::Vector3d, ICPVertex>
    {
    public:
        EIGEN_MAKE_ALIGNED_OPERATOR_NEW
            ICPEdge(Eigen::Vector3d x) : BaseUnaryEdge(), _p(x) { }

        virtual void computeError() override {
            const ICPVertex* v = static_cast<const ICPVertex*>(_vertices[0]);
            const Sophus::SE3d T = v->estimate();
            _error = T * _p - _measurement;
        }

        virtual void linearizeOplus() override {
            const ICPVertex* v = static_cast<const ICPVertex*>(_vertices[0]);
            const Sophus::SE3d T = v->estimate();
            //  e =Tp-p  ,de/T的李代数
            _jacobianOplusXi.block<3, 3>(0, 0) = Eigen::Matrix3d::Identity();
            _jacobianOplusXi.block<3, 3>(0, 3) = -(T.matrix()).block<3, 3>(0, 0) * Sophus::SO3d::hat(_p); //  右乘扰动
            // _jacobianOplusXi.block<3, 3>(0, 3) = -Sophus::SO3d::hat(T * _p); //  左乘扰动
        }

        virtual bool read(std::istream& in) { }
        virtual bool write(std::ostream& out) const { }

    private:
        Eigen::Vector3d _p;
    };

    opt_ICP_G2O::opt_ICP_G2O(const YAML::Node& node)
        : kdtree_flann(new pcl::KdTreeFLANN<PointType>) {

        max_iterations = node["max_iter"].as<int>();
        max_coresspoind_dis = node["max_corr_dist"].as<float>();
        trans_eps = node["trans_eps"].as<float>();
        euc_fitness_eps = node["euc_fitness_eps"].as<float>();
    }

    opt_ICP_G2O::~opt_ICP_G2O() { }

    void opt_ICP_G2O::setInputTarget(const CloudPtr& target) {
        target_ptr = target;
        kdtree_flann->setInputCloud(target);
    }

    void opt_ICP_G2O::setInputSource(const CloudPtr& source) {
        source_ptr = source;
    }

    bool opt_ICP_G2O::align( CloudType& output_ptr, const Eigen::Matrix4f& predict_pose) {
        CloudPtr transform_cloud(new CloudType());
        Eigen::Matrix4f T = predict_pose;
        for (int i = 0; i < max_iterations; ++i) {
            pcl::transformPointCloud(*source_ptr, *transform_cloud, T);

            typedef g2o::BlockSolver<g2o::BlockSolverTraits<6, 3>> BlockSolverType;
            typedef g2o::LinearSolverDense<BlockSolverType::PoseMatrixType> LinearSolverType;

            // auto solver = new g2o::OptimizationAlgorithmGaussNewton(
            //     g2o::make_unique<BlockSolverType>(g2o::make_unique<LinearSolverType>()));

            // typedef LinearSolverCSparse<BlockSolverType::PoseMatrixType> LinearSolverType;

            LinearSolverType* linearSolver = new LinearSolverType();
            // linearSolver->setBlockOrdering(false);
            BlockSolverType* blockSolver = new BlockSolverType(linearSolver);
            auto solver = new g2o::OptimizationAlgorithmGaussNewton(blockSolver);
            // auto solver = new g2o::OptimizationAlgorithmLevenberg(blockSolver);

            g2o::SparseOptimizer opt;
            opt.setAlgorithm(solver);
            opt.setVerbose(false); //    打印
            //  增加顶点
            ICPVertex* v = new ICPVertex();
            v->setEstimate(Sophus::SE3d());
            v->setId(0);
            opt.addVertex(v);

            //  增加边
            for (int j = 0; j < transform_cloud->size(); ++j) {
                const PointType& origin_pt = source_ptr->points[j];
                if (!pcl::isFinite(origin_pt))
                    continue;

                const PointType& transform_pt = transform_cloud->at(j);
                std::vector<float> res_dis;
                std::vector<int> indices;
                kdtree_flann->nearestKSearch(transform_pt, 1, indices, res_dis);
                if (res_dis.front() > max_coresspoind_dis)
                    continue;

                Eigen::Vector3d nearest_pt = Eigen::Vector3d(target_ptr->at(indices.front()).x,
                    target_ptr->at(indices.front()).y,
                    target_ptr->at(indices.front()).z);

                // Eigen::Vector3d origin_eigen(origin_pt.x, origin_pt.y, origin_pt.z);
                Eigen::Vector3d point_eigen(transform_pt.x, transform_pt.y, transform_pt.z);

                //    加入边
                ICPEdge* edge = new ICPEdge(point_eigen);
                edge->setId(j);
                edge->setVertex(0, v);
                edge->setMeasurement(nearest_pt);
                edge->setInformation(Eigen::Matrix<double, 3, 3>::Identity());
                opt.addEdge(edge);
            }

            opt.initializeOptimization();
            opt.optimize(10);
            Sophus::SE3f T_se3 = v->estimate().cast<float>();
            // std::cout << T_se3.matrix() << std::endl;

            T = T_se3.matrix() * T;
        }

        final_pose = T;
        pcl::transformPointCloud(*source_ptr, output_ptr, final_pose);
        return true;
    }

    float opt_ICP_G2O::getFitnessScore() {
        float max_range = std::numeric_limits<float>::max();
        float score = 0.f;

        CloudPtr transform_cloud(new CloudType());
        pcl::transformPointCloud(*source_ptr, *transform_cloud, final_pose);
        std::vector<int> nn_indices(1);
        std::vector<float> nn_dists(1);

        int nr = 0;

        for (size_t i = 0; i < transform_cloud->size(); ++i) {
            kdtree_flann->nearestKSearch(transform_cloud->points[i], 1, nn_indices, nn_dists);
            if (nn_dists.front() <= max_range) {
                score += nn_dists.front();
                nr++;
            }
        }
        if (nr > 0)
            return score / static_cast<float>(nr);
        else
            return (std::numeric_limits<float>::max());
    }
}